package ollama

import (
	"bytes"
	"context"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"io"
	"net/http"
	"net/url"
	"ollama-desktop/internal/config"
	"ollama-desktop/internal/ollama"
	"runtime"
	"strconv"
	"strings"
)

type Client struct {
	Base *url.URL
	Http *http.Client
}

func checkError(resp *http.Response, body []byte) error {
	if resp.StatusCode < http.StatusBadRequest {
		return nil
	}

	return ollama.StatusError{StatusCode: resp.StatusCode, Status: resp.Status, ErrorMessage: string(body)}
}

func (c *Client) do(ctx context.Context, path string, reqData url.Values) ([]byte, error) {
	requestURL := c.Base.JoinPath(path)
	if reqData != nil && len(reqData) > 0 {
		//rawQuery := ""
		//for name, value := range reqData {
		//	if rawQuery != "" {
		//		rawQuery += "&"
		//	}
		//	rawQuery += name + "=" + url.QueryEscape(value)
		//}
		requestURL.RawQuery = reqData.Encode()
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL.String(), nil)
	if err != nil {
		return nil, err
	}

	request.Header.Set("User-Agent", fmt.Sprintf("ollama-desktop/%s (%s %s) Go/%s", config.BuildVersion, runtime.GOARCH, runtime.GOOS, runtime.Version()))

	respObj, err := c.Http.Do(request)
	if err != nil {
		return nil, err
	}
	defer respObj.Body.Close()

	respBody, err := io.ReadAll(respObj.Body)
	if err != nil {
		return nil, err
	}

	if err := checkError(respObj, respBody); err != nil {
		return nil, err
	}
	return respBody, nil
}

// 解析模型名称、是否归档 新界面不存在，暂时不处理
//func parseNameArchive(s *goquery.Selection) (string, bool) {
//	children := s.Children()
//	name := strings.TrimSpace(children.Eq(0).Text())
//	if children.Size() < 2 {
//		return name, false
//	}
//	// 存在名称后面存在标签的情况，比如falcon
//	for i := 1; i < children.Size(); i++ {
//		if strings.Contains(children.Eq(i).Text(), "Archive") {
//			return name, true
//		}
//	}
//	return name, false
//}

func parseName(s *goquery.Selection) string {
	children := s.Children()
	name := strings.TrimSpace(children.Eq(0).Text())
	return name
}

// 解析下载次数、标签数、更新时间
func parsePullsTagsAndUpdated(spans *goquery.Selection) (string, int, string) {
	pullCount := ""
	tagCount := 0
	updated := ""
	spans.Each(func(_ int, info *goquery.Selection) {
		text := strings.ReplaceAll(strings.TrimSpace(info.Text()), "\t", " ")
		if strings.HasSuffix(text, "Pulls") {
			pullCount = strings.TrimSpace(text[:len(text)-5])
		} else if strings.HasSuffix(text, "Pull") {
			pullCount = strings.TrimSpace(text[:len(text)-4])
		} else if strings.HasSuffix(text, "Downloads") {
			pullCount = strings.TrimSpace(text[:len(text)-9])
		} else if strings.HasSuffix(text, "Tag") {
			tagCount, _ = strconv.Atoi(strings.TrimSpace(text[:len(text)-3]))
		} else if strings.HasSuffix(text, "Tags") {
			tagCount, _ = strconv.Atoi(strings.TrimSpace(text[:len(text)-4]))
		} else if strings.HasPrefix(text, "Updated") {
			updated = strings.TrimSpace(text[7:])
		}
	})
	return pullCount, tagCount, updated
}

// 解析分类或者参数大小信息
func parseCategoryOrSize(spans *goquery.Selection) []string {
	var tags []string
	spans.Each(func(_ int, tag *goquery.Selection) {
		tags = append(tags, strings.TrimSpace(tag.Text()))
	})
	return tags
}

func (c *Client) Search(ctx context.Context, request *ollama.SearchRequest) ([]*ollama.SearchItem, error) {
	value := url.Values{}
	value.Set("q", request.Q)
	value.Set("o", request.O)
	for _, item := range request.C {
		value.Add("c", item)
	}
	respBody, err := c.do(ctx, "/search", value)
	if err != nil {
		return nil, err
	}
	doc, err := goquery.NewDocumentFromReader(bytes.NewReader(respBody))
	if err != nil {
		return nil, err
	}
	var list []*ollama.SearchItem

	// 解析模型信息
	doc.Find("ul.grid > li > a").Each(func(_ int, item *goquery.Selection) {
		name := parseName(item.Find("div.flex > h2").First())
		description := strings.TrimSpace(item.Find("div.flex > p.break-words").First().Text())
		categoryOrSize := parseCategoryOrSize(item.Find("div.flex > div.space-x-2 > span"))
		pulls, tags, updated := parsePullsTagsAndUpdated(item.Find("div.flex > p.space-x-5 > span"))
		list = append(list, &ollama.SearchItem{
			Name:           name,
			Description:    description,
			CategoryOrSize: categoryOrSize,
			Pulls:          pulls,
			Tags:           tags,
			Updated:        updated,
		})
	})

	return list, err
}

func (c *Client) Readme(ctx context.Context, modelTag string) (string, error) {
	path := ""
	if strings.Contains(modelTag, "/") {
		path = "/" + modelTag
	} else {
		path = fmt.Sprintf("/library/%s", modelTag)
	}
	respBody, err := c.do(ctx, path, nil)
	if err != nil {
		return "", err
	}
	doc, err := goquery.NewDocumentFromReader(bytes.NewReader(respBody))
	if err != nil {
		return "", err
	}

	readme := doc.Find("div#textareaInput > textarea#editor").Eq(0).Text()

	return readme, err
}

func (c *Client) ModelTags(ctx context.Context, model string) ([]*ollama.ModelTag, error) {
	path := ""
	if strings.Contains(model, "/") {
		path = fmt.Sprintf("/%s/tags", model)
	} else {
		path = fmt.Sprintf("/library/%s/tags", model)
	}
	respBody, err := c.do(ctx, path, nil)
	if err != nil {
		return nil, err
	}
	doc, err := goquery.NewDocumentFromReader(bytes.NewReader(respBody))
	if err != nil {
		return nil, err
	}

	tagsNode := doc.Find("section > div > div > div.px-4.py-3")

	if tagsNode.Size() <= 1 {
		return nil, nil
	}

	var modelTags []*ollama.ModelTag
	for i := 1; i < tagsNode.Size(); i++ {
		tagNode := tagsNode.Eq(i)

		line1Node := tagNode.Find("div > div.grid").First()
		children := line1Node.Children()
		if children.Size() != 4 {
			continue
		}
		// 包含TAG的全名称
		name := strings.TrimSpace(children.Eq(0).Find("span > a").First().Text())
		latest := "latest" == strings.TrimSpace(children.Eq(0).Find("span > a + span").Text())
		size := strings.TrimSpace(children.Eq(1).Text())
		contextSize := strings.TrimSpace(children.Eq(2).Text())
		input := strings.TrimSpace(children.Eq(3).Text())

		infos := strings.Split(tagNode.Find("div > div.grid + div.flex").Text(), "·")
		tag := &ollama.ModelTag{
			Name:    name,
			Latest:  latest,
			Size:    size,
			Context: contextSize,
			Input:   input,
		}
		if len(infos) == 2 {
			tag.Id = strings.TrimSpace(infos[0])
			tag.Updated = strings.TrimSpace(infos[1])
		}
		modelTags = append(modelTags, tag)
	}
	return modelTags, err
}
